import numpy as np
import scipy.special as sc
import scipy.optimize as op
from chaco.api import *
from traits.api import *
from traitsui.api import *
from chaco.tools.api import LineInspector, PanTool, ZoomTool
from enable.api import ComponentEditor


def s21(epsilon, mu, frequency, d):
    k = np.sqrt(mu * epsilon)
    z = np.sqrt(mu / epsilon)
    kd0 = 2 * np.pi * frequency * d / 300
    kd = kd0 * k
    coskd = np.cos(kd)
    jsinkd = 1j * np.sin(kd)
    return 2 / (2 * coskd + jsinkd * (z + 1 / z)) * np.exp(1j * kd0)


class SlabS11S21(HasTraits):
    frequencies = np.arange(.1, 100.1, .1)
    epsilonReal = Range(low=0, high=10.0, value=4.0)
    epsilonImage = Range(low=0, high=10.0, value=.1)
    muReal = Range(low=0, high=10.0, value=2.0)
    muImage = Range(low=0, high=10.0, value=.2)
    thickness = Range(low=0, high=50.0, value=2.0)  # unit mm
    plot = Instance(Plot)
    plotEpsilonMu = Instance(Plot)
    button = Button
    percent = Str('0%')

    def _button_fired(self, event=None):
        epsilonRoot = np.zeros_like(self.frequencies, dtype=complex)
        muRoot = np.zeros_like(self.frequencies, dtype=complex)
        for i, frequency in enumerate(self.frequencies):
            def f(x):
                k = np.sqrt(x[1] * x[0])
                z = np.sqrt(x[1] / x[0])
                coskd = np.cos(2 * np.pi * frequency * self.thickness / 300 * k)
                jsinkd = 1j * np.sin(2 * np.pi * frequency * self.thickness / 300 * k)
                denominator = 2 * coskd + jsinkd * (z + 1 / z)
                s11error = jsinkd * (z - 1 / z) / denominator - self.s11[i]
                s21error = 2 / denominator - self.s21[i]
                return [s11error, s21error]

            r = op.root(f, [4, 1], method='broyden1')
            epsilonRoot[i] = r.x[0]
            muRoot[i] = r.x[1]
            self.percent = f'{i / len(self.frequencies) * 100}%'
        self.plotEpsilonMu.data.update_data(
            rootepsilonReal=epsilonRoot.real,
            rootepsilonImage=-epsilonRoot.imag,
            rootmuReal=muRoot.real,
            rootmuImage=-muRoot.imag,
        )
        for i, frequency in enumerate(self.frequencies):
            print(i, frequency, epsilonRoot[i], muRoot[i])

    def traits_init(self):
        self.plot = Plot(ArrayPlotData())
        self.plotEpsilonMu = Plot(ArrayPlotData())
        self.plot.data.update_data(
            frequencies=np.array([]),
            s11=np.array([]),
            s21=np.array([]),
        )
        self.plotEpsilonMu.data.update_data(
            frequencies=np.array([]),
            epsilonReal=np.array([]),
            epsilonImage=np.array([]),
            muReal=np.array([]),
            muImage=np.array([]),
            nrwepsilonReal=np.array([]),
            nrwepsilonImage=np.array([]),
            nrwmuReal=np.array([]),
            nrwmuImage=np.array([]),
            rootepsilonReal=np.array([]),
            rootepsilonImage=np.array([]),
            rootmuReal=np.array([]),
            rootmuImage=np.array([]),
        )
        self.update_data()
        self.plot.title = 'S11 and S21'
        self.plot.x_axis.title = "Frequency / GHz"
        self.plot.y_axis.title = "S11 and S21"
        self.plot.plot(("frequencies", "s11"), name='S11', color=palette11[0])
        self.plot.plot(("frequencies", "s21"), name='S21', color=palette11[1])
        legend = Legend()
        legend.plots = self.plot.plots
        self.plot.overlays.append(legend)
        self.plotEpsilonMu.title = 'Epsilon and Mu'
        self.plotEpsilonMu.x_axis.title = "Frequency / GHz"
        self.plotEpsilonMu.y_axis.title = "Epsilon and Mu"
        self.plotEpsilonMu.plot(("frequencies", "epsilonReal"), name='Epsilon Real', color=palette11[0], line_style="dot")
        self.plotEpsilonMu.plot(("frequencies", "epsilonImage"), name='Epsilon Image', color=palette11[1], line_style="dot")
        self.plotEpsilonMu.plot(("frequencies", "muReal"), name='Mu Real', color=palette11[2], line_style="dot")
        self.plotEpsilonMu.plot(("frequencies", "muImage"), name='Mu Image', color=palette11[3], line_style="dot")
        self.plotEpsilonMu.plot(("frequencies", "nrwepsilonReal"), name='Epsilon Real NRW', color=palette11[0], line_width=3)
        self.plotEpsilonMu.plot(("frequencies", "nrwepsilonImage"), name='Epsilon Image NRW', color=palette11[1], line_width=3)
        self.plotEpsilonMu.plot(("frequencies", "nrwmuReal"), name='Mu Real NRW', color=palette11[2], line_width=3)
        self.plotEpsilonMu.plot(("frequencies", "nrwmuImage"), name='Mu Image NRW', color=palette11[3], line_width=3)
        self.plotEpsilonMu.plot(("frequencies", "rootepsilonReal"), name='Epsilon Real Root', color=palette11[0], line_width=3, line_style="dot")
        self.plotEpsilonMu.plot(("frequencies", "rootepsilonImage"), name='Epsilon Image Root', color=palette11[1], line_width=3, line_style="dot")
        self.plotEpsilonMu.plot(("frequencies", "rootmuReal"), name='Mu Real Root', color=palette11[2], line_width=3, line_style="dot")
        self.plotEpsilonMu.plot(("frequencies", "rootmuImage"), name='Mu Image Root', color=palette11[3], line_width=3, line_style="dot")
        self.plotEpsilonMu.range2d.set_bounds((0, -20), (self.frequencies[-1], 20))
        legend = Legend()
        legend.plots = self.plotEpsilonMu.plots
        self.plotEpsilonMu.overlays.append(legend)
        self.plotEpsilonMu.tools.append(
            PanTool(self.plotEpsilonMu, constrain_key="shift")
        )
        self.plotEpsilonMu.overlays.append(
            ZoomTool(component=self.plotEpsilonMu, tool_mode="box", always_on=False)
        )

    @observe('epsilonReal,epsilonImage,muReal,muImage,thickness')
    def update_data(self, event=None):
        epsilon = self.epsilonReal - 1j * self.epsilonImage
        mu = self.muReal - 1j * self.muImage
        k = np.sqrt(epsilon * mu)
        z = mu / k
        coskd = np.cos(2 * np.pi * self.frequencies * self.thickness / 300 * k)
        jsinkd = 1j * np.sin(2 * np.pi * self.frequencies * self.thickness / 300 * k)
        denominator = 2 * coskd + jsinkd * (z + 1 / z)
        self.s11 = jsinkd * (z - 1 / z) / denominator
        self.s21 = 2 / denominator
        self.plot.data.update_data(
            frequencies=self.frequencies,
            s11=20 * np.log10(np.abs(self.s11)),
            s21=20 * np.log10(np.abs(self.s21)),
        )
        zNRW = ((1 + self.s11) ** 2 - self.s21 ** 2) / ((1 - self.s11) ** 2 - self.s21 ** 2)
        zNRW = np.sqrt(zNRW)
        kNRW = ((1 - self.s21) ** 2 - self.s11 ** 2) / ((1 + self.s21) ** 2 - self.s11 ** 2)
        kNRW = np.arctan(-1j * np.sqrt(kNRW)) * 600 / 2 / np.pi / self.frequencies / self.thickness
        epsilonNRW = kNRW / zNRW
        muNRW = kNRW * zNRW
        self.plotEpsilonMu.data.update_data(
            frequencies=self.frequencies,
            epsilonReal=np.full_like(self.frequencies, self.epsilonReal),
            epsilonImage=np.full_like(self.frequencies, self.epsilonImage),
            muReal=np.full_like(self.frequencies, self.muReal),
            muImage=np.full_like(self.frequencies, self.muImage),
            nrwepsilonReal=epsilonNRW.real,
            nrwepsilonImage=-epsilonNRW.imag,
            nrwmuReal=muNRW.real,
            nrwmuImage=-muNRW.imag,
        )

    traits_view = View(
        Tabbed(
            UItem(
                "plot",
                editor=ComponentEditor(),
                resizable=True,
                label='S参数'
            ),
            UItem(
                "plotEpsilonMu",
                editor=ComponentEditor(),
                resizable=True,
                label='电磁参数'

            )
        ),
        Item('epsilonReal', label='介电常数实部'),
        Item('epsilonImage', label='介电常数虚部'),
        Item('muReal', label='磁导率实部'),
        Item('muImage', label='磁导率虚部'),
        Item('thickness', label='厚度/mm'),
        HGroup(
            UItem('button', label='计算'),
            Item('percent', label='进度', style='readonly'),
            spring
        ),
        width=1400,
        height=1000,
        resizable=True,
    )


class TwoSlabS21(HasTraits):
    frequencies = np.arange(.1, 100.1, .1)
    epsilonReal = Range(low=0, high=10.0, value=4.0)
    epsilonImage = Range(low=0, high=10.0, value=.1)
    muReal = Range(low=0, high=10.0, value=2.0)
    muImage = Range(low=0, high=10.0, value=.2)
    thickness1 = Range(low=0, high=50.0, value=2.0)  # unit mm
    thickness2 = Range(low=0, high=50.0, value=3.0)  # unit mm
    plot = Instance(Plot)
    plotEpsilonMu = Instance(Plot)
    button = Button
    percent = Str('0%')
    initialEpsilon = Complex(4 - 1j)
    initialMu = Complex(2 - 1j)

    def _button_fired(self, event=None):
        epsilonRoot = np.zeros_like(self.frequencies, dtype=complex)
        muRoot = np.zeros_like(self.frequencies, dtype=complex)
        for i, frequency in enumerate(self.frequencies):
            def f(x):
                return [
                    s21(x[0], x[1], frequency, self.thickness1) - self.slab1s21[i],
                    s21(x[0], x[1], frequency, self.thickness2) - self.slab2s21[i],
                ]
            if i == 0:
                r = op.root(f, [self.initialEpsilon, self.initialMu], method='broyden1')
            else:
                r = op.root(f, [epsilonRoot[i - 1], muRoot[i - 1]], method='broyden1')
            epsilonRoot[i] = r.x[0]
            muRoot[i] = r.x[1]
            print(i, frequency, epsilonRoot[i], muRoot[i])
            self.plotEpsilonMu.data.update_data(
                frequencies=self.frequencies,
                rootepsilonReal=epsilonRoot.real,
                rootepsilonImage=-epsilonRoot.imag,
                rootmuReal=muRoot.real,
                rootmuImage=-muRoot.imag,
            )

    def traits_init(self):
        self.plot = Plot(ArrayPlotData())
        self.plot.data.update_data(
            frequencies=np.array([]),
            slab1s21=np.array([]),
            slab2s21=np.array([]),
            rootslab1s21=np.array([]),
            rootslab2s21=np.array([]),
        )
        self.plot.title = 'S21'
        self.plot.x_axis.title = "Frequency / GHz"
        self.plot.y_axis.title = "S21"
        self.plot.plot(("frequencies", "slab1s21"), name='slab1 S21', color=palette11[0])
        self.plot.plot(("frequencies", "slab2s21"), name='slab2 S21', color=palette11[1])
        self.plot.plot(("frequencies", "rootslab1s21"), name='slab1 S21', color=palette11[0],line_width=3,line_style='dot')
        self.plot.plot(("frequencies", "rootslab2s21"), name='slab2 S21', color=palette11[1],line_width=3,line_style='dot')
        legend = Legend()
        legend.plots = self.plot.plots
        self.plot.overlays.append(legend)
        self.plotEpsilonMu = Plot(ArrayPlotData())
        self.plotEpsilonMu.data.update_data(
            frequencies=np.array([]),
            epsilonReal=np.array([]),
            epsilonImage=np.array([]),
            muReal=np.array([]),
            muImage=np.array([]),
            rootepsilonReal=np.array([]),
            rootepsilonImage=np.array([]),
            rootmuReal=np.array([]),
            rootmuImage=np.array([]),
        )
        self.plotEpsilonMu.title = 'Epsilon and Mu'
        self.plotEpsilonMu.x_axis.title = "Frequency / GHz"
        self.plotEpsilonMu.y_axis.title = "Epsilon and Mu"
        self.plotEpsilonMu.plot(("frequencies", "epsilonReal"), name='Epsilon Real', color=palette11[0])
        self.plotEpsilonMu.plot(("frequencies", "epsilonImage"), name='Epsilon Image', color=palette11[1], line_style="dot")
        self.plotEpsilonMu.plot(("frequencies", "muReal"), name='Mu Real', color=palette11[2])
        self.plotEpsilonMu.plot(("frequencies", "muImage"), name='Mu Image', color=palette11[3], line_style="dot")
        self.plotEpsilonMu.plot(("frequencies", "rootepsilonReal"), name='Epsilon Real Root', color=palette11[0], line_width=3)
        self.plotEpsilonMu.plot(("frequencies", "rootepsilonImage"), name='Epsilon Image Root', color=palette11[1], line_width=3, line_style="dot")
        self.plotEpsilonMu.plot(("frequencies", "rootmuReal"), name='Mu Real Root', color=palette11[2], line_width=3)
        self.plotEpsilonMu.plot(("frequencies", "rootmuImage"), name='Mu Image Root', color=palette11[3], line_width=3, line_style="dot")
        self.plotEpsilonMu.range2d.set_bounds((0, -20), (self.frequencies[-1], 20))
        legend = Legend()
        legend.plots = self.plotEpsilonMu.plots
        self.plotEpsilonMu.overlays.append(legend)
        self.plotEpsilonMu.tools.append(
            PanTool(self.plotEpsilonMu, constrain_key="shift")
        )
        self.plotEpsilonMu.overlays.append(
            ZoomTool(component=self.plotEpsilonMu, tool_mode="box", always_on=False)
        )
        self.update_data()

    @observe('epsilonReal,epsilonImage,muReal,muImage,thickness1,thickness2')
    def update_data(self, event=None):
        epsilon = self.epsilonReal - 1j * self.epsilonImage
        mu = self.muReal - 1j * self.muImage
        self.slab1s21 = s21(epsilon,mu,self.frequencies,self.thickness1)
        self.slab2s21 = s21(epsilon,mu,self.frequencies,self.thickness2)
        self.plot.data.update_data(
            frequencies=self.frequencies,
            slab1s21=20 * np.log10(np.abs(self.slab1s21)),
            slab2s21=20 * np.log10(np.abs(self.slab2s21)),
        )
        self.plotEpsilonMu.data.update_data(
            frequencies=self.frequencies,
            epsilonReal=np.full_like(self.frequencies, self.epsilonReal),
            epsilonImage=np.full_like(self.frequencies, self.epsilonImage),
            muReal=np.full_like(self.frequencies, self.muReal),
            muImage=np.full_like(self.frequencies, self.muImage),
        )
    traits_view = View(
        Tabbed(
            UItem(
                "plotEpsilonMu",
                editor=ComponentEditor(),
                resizable=True,
                label='电磁参数'

            ),
            UItem(
                "plot",
                editor=ComponentEditor(),
                resizable=True,
                label='S参数'
            )
        ),
        Item('epsilonReal', label='介电常数实部'),
        Item('epsilonImage', label='介电常数虚部'),
        Item('muReal', label='磁导率实部'),
        Item('muImage', label='磁导率虚部'),
        Item('thickness1', label='介质板1厚度/mm'),
        Item('thickness2', label='介质板2厚度/mm'),
        HGroup(
            UItem('button', label='计算'),
            Item('percent', label='进度', style='readonly'),
            spring
        ),
        width=1400,
        height=1000,
        resizable=True,
    )


class TwoSlabS21Experiment(HasTraits):
    file1 = File('DMF_0.86mm.tdat')
    file2 = File('DMF_1.86mm.tdat')
    fileMeasure = File('DMF_Jijinzu.eudat')
    thickness1 = Float(.86)  # unit mm
    thickness2 = Float(1.78)  # unit mm
    buttonCalculate = Button
    buttonImport = Button
    buttonImportMeasure = Button
    initialEpsilon = Complex(30 - 5j)
    initialMu = Complex(10 - 4j)

    def _buttonImportMeasure_fired(self, event):
        data=np.loadtxt(self.fileMeasure)
        self.plotEpsilonMu.data.update_data(
            measurefrequencies=data[:,0],
            measureepsilonReal=data[:,1],
            measureepsilonImage=data[:,2],
            measuremuReal=data[:,3],
            measuremuImage=data[:,4],
        )

    def _buttonImport_fired(self, event):
        with open(self.file1) as f:
            firstLine = f.readline().split()
            frequencyNumber = int(firstLine[0])
            frequencyStart = float(firstLine[1])
            frequencyStop = float(firstLine[2])
            self.frequencies = np.linspace(frequencyStart, frequencyStop, frequencyNumber)
            a = np.loadtxt(self.file1, skiprows=1, max_rows=frequencyNumber)
            self.slab1s21 = a[:, 0] + 1j * a[:, 1]
        with open(self.file2) as f:
            firstLine = f.readline().split()
            frequencyNumber = int(firstLine[0])
            frequencyStart = float(firstLine[1])
            frequencyStop = float(firstLine[2])
            self.frequencies = np.linspace(frequencyStart, frequencyStop, frequencyNumber)
            a = np.loadtxt(self.file2, skiprows=1, max_rows=frequencyNumber)
            self.slab2s21 = a[:, 0] + 1j * a[:, 1]
        self.plot.data.update_data(
            frequencies=self.frequencies,
            slab1s21=20 * np.log10(np.abs(self.slab1s21)),
            slab2s21=20 * np.log10(np.abs(self.slab2s21)),
        )
        self.plotS21Phase.data.update_data(
            frequencies=self.frequencies,
            slab1s21=np.angle(self.slab1s21, deg=True),
            slab2s21=np.angle(self.slab2s21, deg=True),
        )

    def _buttonCalculate_fired(self, event=None):
        epsilonRoot = np.zeros_like(self.frequencies, dtype=complex)
        muRoot = np.zeros_like(self.frequencies, dtype=complex)
        for i, frequency in enumerate(self.frequencies):
            def f(x):
                return [
                    s21(x[1], x[0], frequency, self.thickness1) - self.slab1s21[i],
                    s21(x[1], x[0], frequency, self.thickness2) - self.slab2s21[i],
                ]

            if i == 0:
                r = op.root(f, [self.initialEpsilon, self.initialMu], method='broyden1')
            else:
                r = op.root(f, [epsilonRoot[i - 1], muRoot[i - 1]], method='broyden1')
            epsilonRoot[i] = r.x[0]
            muRoot[i] = r.x[1]
            print(i, frequency, epsilonRoot[i], muRoot[i])
        self.plotEpsilonMu.data.update_data(
            frequencies=self.frequencies,
            rootepsilonReal=epsilonRoot.real,
            rootepsilonImage=-epsilonRoot.imag,
            rootmuReal=muRoot.real,
            rootmuImage=-muRoot.imag,
        )
        s21_1 = s21(epsilonRoot, muRoot, self.frequencies, self.thickness1)
        s21_2 = s21(epsilonRoot, muRoot, self.frequencies, self.thickness2)
        self.plot.data.update_data(
            frequencies=self.frequencies,
            rootslab1s21=20 * np.log10(np.abs(s21_1)),
            rootslab2s21=20 * np.log10(np.abs(s21_2)),
        )
        self.plotS21Phase.data.update_data(
            frequencies=self.frequencies,
            rootslab1s21=np.angle(s21_1, deg=True),
            rootslab2s21=np.angle(s21_2, deg=True),
        )

    def traits_init(self):
        self.plot = Plot(ArrayPlotData())
        self.plot.data.update_data(
            frequencies=np.array([]),
            slab1s21=np.array([]),
            slab2s21=np.array([]),
            rootslab1s21=np.array([]),
            rootslab2s21=np.array([]),
        )
        self.plot.title = 'S21'
        self.plot.x_axis.title = "Frequency / GHz"
        self.plot.y_axis.title = "S21 magnitude / dB"
        self.plot.plot(("frequencies", "slab1s21"), name='slab1 S21', color=palette11[0])
        self.plot.plot(("frequencies", "slab2s21"), name='slab2 S21', color=palette11[1])
        self.plot.plot(("frequencies", "rootslab1s21"), name='slab1 S21 root', color=palette11[0], line_style='dot', line_width=2)
        self.plot.plot(("frequencies", "rootslab2s21"), name='slab2 S21 root', color=palette11[1], line_style='dot', line_width=2)
        legend = Legend()
        legend.plots = self.plot.plots
        self.plot.overlays.append(legend)
        self.plot.tools.append(
            PanTool(self.plot, constrain_key="shift")
        )
        self.plot.overlays.append(
            ZoomTool(component=self.plot, tool_mode="box", always_on=False)
        )

        self.plotS21Phase = Plot(ArrayPlotData())
        self.plotS21Phase.data.update_data(
            frequencies=np.array([]),
            slab1s21=np.array([]),
            slab2s21=np.array([]),
            rootslab1s21=np.array([]),
            rootslab2s21=np.array([]),
        )
        self.plotS21Phase.title = 'Phase'
        self.plotS21Phase.x_axis.title = "Frequency / GHz"
        self.plotS21Phase.y_axis.title = "Phase / deg"
        self.plotS21Phase.plot(("frequencies", "slab1s21"), name='slab1 S21', color=palette11[0])
        self.plotS21Phase.plot(("frequencies", "slab2s21"), name='slab2 S21', color=palette11[1])
        self.plotS21Phase.plot(("frequencies", "rootslab1s21"), name='slab1 S21 root', color=palette11[0], line_style='dot', line_width=2)
        self.plotS21Phase.plot(("frequencies", "rootslab2s21"), name='slab2 S21 root', color=palette11[1], line_style='dot', line_width=2)
        legend = Legend()
        legend.plots = self.plotS21Phase.plots
        self.plotS21Phase.overlays.append(legend)
        self.plotS21Phase.tools.append(
            PanTool(self.plotS21Phase, constrain_key="shift")
        )
        self.plotS21Phase.overlays.append(
            ZoomTool(component=self.plotS21Phase, tool_mode="box", always_on=False)
        )

        self.plotEpsilonMu = Plot(ArrayPlotData())
        self.plotEpsilonMu.data.update_data(
            frequencies=np.array([]),
            rootepsilonReal=np.array([]),
            rootepsilonImage=np.array([]),
            rootmuReal=np.array([]),
            rootmuImage=np.array([]),
            measurefrequencies=np.array([]),
            measureepsilonReal=np.array([]),
            measureepsilonImage=np.array([]),
            measuremuReal=np.array([]),
            measuremuImage=np.array([]),
        )
        self.plotEpsilonMu.title = 'Epsilon and Mu'
        self.plotEpsilonMu.x_axis.title = "Frequency / GHz"
        self.plotEpsilonMu.y_axis.title = "Epsilon and Mu"
        self.plotEpsilonMu.plot(("frequencies", "rootepsilonReal"), name='Epsilon Real Root', color='blue', line_width=2)
        self.plotEpsilonMu.plot(("frequencies", "rootepsilonImage"), name='Epsilon Image Root', color='blue', line_width=2, line_style='dot')
        self.plotEpsilonMu.plot(("frequencies", "rootmuReal"), name='Mu Real Root', color='red', line_width=2)
        self.plotEpsilonMu.plot(("frequencies", "rootmuImage"), name='Mu Image Root', color='red', line_width=2, line_style='dot')
        self.plotEpsilonMu.plot(("measurefrequencies", "measureepsilonReal"), name='Epsilon Real Measure', color='blue', line_width=4)
        self.plotEpsilonMu.plot(("measurefrequencies", "measureepsilonImage"), name='Epsilon Image Measure', color='blue', line_width=4, line_style='dot')
        self.plotEpsilonMu.plot(("measurefrequencies", "measuremuReal"), name='Mu Real Measure', color='red', line_width=4)
        self.plotEpsilonMu.plot(("measurefrequencies", "measuremuImage"), name='Mu Image Measure', color='red', line_width=4, line_style='dot')
        legend = Legend()
        legend.plots = self.plotEpsilonMu.plots
        self.plotEpsilonMu.overlays.append(legend)
        self.plotEpsilonMu.tools.append(
            PanTool(self.plotEpsilonMu, constrain_key="shift")
        )
        self.plotEpsilonMu.overlays.append(
            ZoomTool(component=self.plotEpsilonMu, tool_mode="box", always_on=False)
        )

    traits_view = View(
        Tabbed(
            Group(
                Tabbed(
                    UItem(
                        "plot",
                        editor=ComponentEditor(),
                        resizable=True,
                        label='S参数幅度'
                    ),
                    UItem(
                        "plotS21Phase",
                        editor=ComponentEditor(),
                        resizable=True,
                        label='S参数相位'
                    )
                ),
                HGroup(
                    Item('thickness1', label='介质板1厚度/mm'),
                    Item('file1', label='厚度1测试结果'),
                ),
                HGroup(
                    Item('thickness2', label='介质板2厚度/mm'),
                    Item('file2', label='厚度2测试结果'),
                ),
                UItem('buttonImport', label='导入'),
                label='透射率'
            ),
            Group(
                UItem(
                    "plotEpsilonMu",
                    editor=ComponentEditor(),
                    resizable=True,
                    label='电磁参数'
                ),
                HGroup(
                    UItem('buttonCalculate', label='反演'),
                    UItem('buttonImportMeasure', label='导入测试结果'),
                    Item('initialEpsilon', label='介电常数初值'),
                    Item('initialMu', label='磁导率初值'),
                ),
                Item('fileMeasure', label='测试结果'),
                label='电磁参数'
            ),
        ),
        width=1400,
        height=1000,
        resizable=True,
    )


class OneSlabS21Experiment(HasTraits):
    file = File('E:/pip310/Dielectic__0.25mm.tdat')
    thickness = Float(.25)  # unit mm
    buttonImport = Button
    buttonCalculate = Button
    initialEpsilon = Complex(30 - 5j)
    initialMu = Complex(10 - 4j)
    asignment = Enum('指定介电常数', '指定磁导率', )
    asignedEpsilon = Complex(30 - 5j)
    asignedMu = Complex(1)

    def _buttonImport_fired(self, event):
        with open(self.file) as f:
            firstLine = f.readline().split()
            frequencyNumber = int(firstLine[0])
            frequencyStart = float(firstLine[1])
            frequencyStop = float(firstLine[2])
            self.frequencies = np.linspace(frequencyStart, frequencyStop, frequencyNumber)
            a = np.loadtxt(self.file, skiprows=1, max_rows=frequencyNumber)
            self.s21 = a[:, 0] + 1j * a[:, 1]
        self.plot.data.update_data(
            frequencies=self.frequencies,
            s21=20 * np.log10(np.abs(self.s21)),
        )
        self.plotS21Phase.data.update_data(
            frequencies=self.frequencies,
            s21=np.angle(self.s21, deg=True),
        )

    def _buttonCalculate_fired(self, event=None):
        epsilonmuRoot = np.zeros_like(self.frequencies, dtype=complex)
        if self.asignment == '指定介电常数':
            for i, frequency in enumerate(self.frequencies):
                def f(x):
                    return s21(self.asignedEpsilon, x, frequency, self.thickness) - self.s21[i]

                if i == 0:
                    r = op.root(f, self.initialMu, method='broyden1')
                else:
                    r = op.root(f, epsilonmuRoot[i - 1], method='broyden1')
                epsilonmuRoot[i] = r.x[1]
                print(i, frequency, epsilonmuRoot[i])
            s21_recalculate = s21(self.asignedEpsilon, epsilonmuRoot, self.frequencies, self.thickness)
        elif self.asignment == '指定磁导率':
            for i, frequency in enumerate(self.frequencies):
                def f(x):
                    return s21(x, self.asignedMu, frequency, self.thickness) - self.s21[i]

                if i == 0:
                    r = op.root(f, self.initialEpsilon, method='broyden1')
                else:
                    r = op.root(f, epsilonmuRoot[i - 1], method='broyden1')
                epsilonmuRoot[i] = r.x
                print(i, frequency, epsilonmuRoot[i])
            s21_recalculate = s21(epsilonmuRoot, self.asignedMu, self.frequencies, self.thickness)
        self.plotEpsilonMu.data.update_data(
            frequencies=self.frequencies,
            rootepsilonmuReal=epsilonmuRoot.real,
            rootepsilonmuImage=-epsilonmuRoot.imag,
        )
        self.plot.data.update_data(
            frequencies=self.frequencies,
            roots21=20 * np.log10(np.abs(s21_recalculate)),
        )
        self.plotS21Phase.data.update_data(
            frequencies=self.frequencies,
            roots21=np.angle(s21_recalculate, deg=True),
        )

    def traits_init(self):
        self.asignment = '指定磁导率'
        self.plot = Plot(ArrayPlotData())
        self.plot.data.update_data(
            frequencies=np.array([]),
            s21=np.array([]),
            roots21=np.array([]),
        )
        self.plot.title = 'S21'
        self.plot.x_axis.title = "Frequency / GHz"
        self.plot.y_axis.title = "S21 magnitude / dB"
        self.plot.plot(("frequencies", "s21"), name='S21', color=palette11[0])
        self.plot.plot(("frequencies", "roots21"), name='S21 root', color=palette11[0], line_style='dot', line_width=2)
        legend = Legend()
        legend.plots = self.plot.plots
        self.plot.overlays.append(legend)
        self.plot.tools.append(
            PanTool(self.plot, constrain_key="shift")
        )
        self.plot.overlays.append(
            ZoomTool(component=self.plot, tool_mode="box", always_on=False)
        )

        self.plotS21Phase = Plot(ArrayPlotData())
        self.plotS21Phase.data.update_data(
            frequencies=np.array([]),
            s21=np.array([]),
            roots21=np.array([]),
        )
        self.plotS21Phase.title = 'Phase'
        self.plotS21Phase.x_axis.title = "Frequency / GHz"
        self.plotS21Phase.y_axis.title = "Phase / deg"
        self.plotS21Phase.plot(("frequencies", "s21"), name='S21', color=palette11[0])
        self.plotS21Phase.plot(("frequencies", "roots21"), name='S21 root', color=palette11[0], line_style='dot', line_width=2)
        legend = Legend()
        legend.plots = self.plotS21Phase.plots
        self.plotS21Phase.overlays.append(legend)
        self.plotS21Phase.tools.append(
            PanTool(self.plotS21Phase, constrain_key="shift")
        )
        self.plotS21Phase.overlays.append(
            ZoomTool(component=self.plotS21Phase, tool_mode="box", always_on=False)
        )

        self.plotEpsilonMu = Plot(ArrayPlotData())
        self.plotEpsilonMu.data.update_data(
            frequencies=np.array([]),
            rootepsilonmuReal=np.array([]),
            rootepsilonmuImage=np.array([]),
        )
        self.plotEpsilonMu.title = 'Epsilon and Mu'
        self.plotEpsilonMu.x_axis.title = "Frequency / GHz"
        self.plotEpsilonMu.y_axis.title = "Epsilon and Mu"
        self.plotEpsilonMu.plot(("frequencies", "rootepsilonmuReal"), name='Real Root', color='blue', line_width=2)
        self.plotEpsilonMu.plot(("frequencies", "rootepsilonmuImage"), name='Image Root', color='blue', line_width=2, line_style='dot')
        # self.plotEpsilonMu.range2d.set_bounds((2, -5), (18, 20))
        legend = Legend()
        legend.plots = self.plotEpsilonMu.plots
        self.plotEpsilonMu.overlays.append(legend)
        self.plotEpsilonMu.tools.append(
            PanTool(self.plotEpsilonMu, constrain_key="shift")
        )
        self.plotEpsilonMu.overlays.append(
            ZoomTool(component=self.plotEpsilonMu, tool_mode="box", always_on=False)
        )

    traits_view = View(
        Tabbed(
            Group(
                Tabbed(
                    UItem(
                        "plot",
                        editor=ComponentEditor(),
                        resizable=True,
                        label='S参数幅度'
                    ),
                    UItem(
                        "plotS21Phase",
                        editor=ComponentEditor(),
                        resizable=True,
                        label='S参数相位'
                    )
                ),
                HGroup(
                    UItem('buttonImport', label='导入'),
                    Item('thickness', label='介质板厚度/mm', width=300),
                    Item('file', label='透波率测试结果'),
                ),
                label='透射率'
            ),
            Group(
                UItem(
                    "plotEpsilonMu",
                    editor=ComponentEditor(),
                    resizable=True,
                    label='电磁参数'
                ),
                HGroup(
                    UItem('buttonCalculate', label='计算'),
                    Item('asignment', label='指定介质', style='custom'),
                    HGroup(
                        Item('asignedEpsilon', label='介电常数', enabled_when='asignment=="指定介电常数"'),
                        Item('asignedMu', label='磁导率', enabled_when='asignment=="指定磁导率"'),
                        label='指定参数',
                        show_border=True,
                    ),
                    HGroup(
                        Item('initialEpsilon', label='介电常数', enabled_when='asignment=="指定磁导率"'),
                        Item('initialMu', label='磁导率', enabled_when='asignment=="指定介电常数"'),
                        label='参数初始值',
                        show_border=True,
                    ),
                ),
                label='电磁参数'
            ),
        ),
        width=1400,
        height=1000,
        resizable=True,
    )


class Form(HasTraits):
    slabS11S21 = Instance(SlabS11S21, ())
    twoSlabS21 = Instance(TwoSlabS21, ())
    twoSlabS21Experiment = Instance(TwoSlabS21Experiment, ())
    oneSlabS21Experiment = Instance(OneSlabS21Experiment, ())
    traits_view = View(
        Tabbed(
            UItem('@twoSlabS21Experiment',label='两透射率反演'),
            UItem('@oneSlabS21Experiment',label='单透射率反演'),
            UItem('@slabS11S21',label='NRW法验证'),
            UItem('@twoSlabS21',label='两透射率验证'),
        ),
        width=1400,
        height=1000,
        resizable=True,
        title='电磁参数反演'
    )

if __name__ == '__main__':
    form = Form()
    form.configure_traits()
